# 多模态增强器
import re
import json
from typing import Dict, List, Any, Optional

class MultimodalEnhancer:
    def __init__(self):
        self.mermaid_counter = 0
    
    def _normalize_markdown_codeblock(self, text: str) -> str:
        """还原被转义的markdown代码块前缀"""
        import re
        # 先还原mermaid和rust等带语言的
        text = re.sub(r'\\+`{3}\s*mermaid', '```mermaid', text)
        text = re.sub(r'\\+`{3}\s*rust', '```rust', text)
        # 再还原普通代码块
        text = re.sub(r'\\+`{3}', '```', text)
        return text

    def enhance_content(self, content: str) -> Dict[str, Any]:
        """增强内容，提取多模态元素"""
        # 先去除markdown代码块前缀的多余转义
        content = self._normalize_markdown_codeblock(content)
        enhanced = {
            'text_content': content,
            'code_blocks': [],
            'mermaid_diagrams': [],
            'tables': [],
            'images': [],
            'enhanced_html': content
        }
        
        # 提取代码块
        enhanced['code_blocks'] = self._extract_code_blocks(content)
        
        # 提取Mermaid图表
        enhanced['mermaid_diagrams'] = self._extract_mermaid_diagrams(content)
        
        # 提取表格
        enhanced['tables'] = self._extract_tables(content)
        
        # 生成增强的HTML
        enhanced['enhanced_html'] = self._generate_enhanced_html(content, enhanced)
        
        return enhanced
    
    def _extract_code_blocks(self, content: str) -> List[Dict[str, str]]:
        """提取代码块"""
        code_blocks = []
        pattern = r'```(\w+)?\n(.*?)```'
        matches = re.finditer(pattern, content, re.DOTALL)
        
        for match in matches:
            language = match.group(1) or 'text'
            code = match.group(2).strip()
            code_blocks.append({
                'language': language,
                'code': code,
                'id': f'code-{len(code_blocks)}'
            })
        
        return code_blocks
    
    def _extract_mermaid_diagrams(self, content: str) -> List[Dict[str, str]]:
        """提取Mermaid图表"""
        mermaid_diagrams = []
        pattern = r'```mermaid\s*\n(.*?)```'
        matches = re.finditer(pattern, content, re.DOTALL)
        
        for match in matches:
            diagram_code = match.group(1).strip()
            self.mermaid_counter += 1
            mermaid_diagrams.append({
                'id': f'mermaid-{self.mermaid_counter}',
                'code': diagram_code,
                'type': self._detect_mermaid_type(diagram_code)
            })
        
        return mermaid_diagrams
    
    def _detect_mermaid_type(self, diagram_code: str) -> str:
        """检测Mermaid图表类型"""
        if 'graph' in diagram_code or 'flowchart' in diagram_code:
            return 'flowchart'
        elif 'sequenceDiagram' in diagram_code:
            return 'sequence'
        elif 'classDiagram' in diagram_code:
            return 'class'
        elif 'stateDiagram' in diagram_code:
            return 'state'
        elif 'gantt' in diagram_code:
            return 'gantt'
        elif 'pie' in diagram_code:
            return 'pie'
        else:
            return 'unknown'
    
    def _extract_tables(self, content: str) -> List[Dict[str, Any]]:
        """提取Markdown表格"""
        tables = []
        pattern = r'\|.*\|\n\|[\s\-:|]+\|\n(\|.*\|\n)*'
        matches = re.finditer(pattern, content)
        
        for match in matches:
            table_text = match.group(0)
            table_data = self._parse_markdown_table(table_text)
            tables.append({
                'id': f'table-{len(tables)}',
                'data': table_data,
                'html': self._table_to_html(table_data)
            })
        
        return tables
    
    def _parse_markdown_table(self, table_text: str) -> List[List[str]]:
        """解析Markdown表格"""
        lines = table_text.strip().split('\n')
        table_data = []
        
        for line in lines:
            if line.startswith('|') and line.endswith('|'):
                # 移除首尾的|并分割
                cells = [cell.strip() for cell in line[1:-1].split('|')]
                table_data.append(cells)
        
        return table_data
    
    def _table_to_html(self, table_data: List[List[str]]) -> str:
        """将表格数据转换为HTML"""
        if not table_data:
            return ""
        
        html = '<table class="enhanced-table">\n'
        
        for i, row in enumerate(table_data):
            if i == 0:
                html += '  <thead>\n    <tr>\n'
                for cell in row:
                    html += f'      <th>{cell}</th>\n'
                html += '    </tr>\n  </thead>\n  <tbody>\n'
            else:
                html += '    <tr>\n'
                for cell in row:
                    html += f'      <td>{cell}</td>\n'
                html += '    </tr>\n'
        
        html += '  </tbody>\n</table>'
        return html
    
    def _generate_enhanced_html(self, content: str, enhanced: Dict[str, Any]) -> str:
        """生成增强的HTML内容（美化代码块和mermaid图表，并分块分段，p标签仅包裹纯文本）"""
        html_content = content

        # 1. 标题美化
        html_content = re.sub(r'^### (.*)$', r'<h3 style="margin: 18px 0 8px 0; font-size: 1.1em; color: #2d3748;">\1</h3>', html_content, flags=re.MULTILINE)
        html_content = re.sub(r'^## (.*)$', r'<h2 style="margin: 28px 0 12px 0; font-size: 1.3em; color: #2563eb;">\1</h2>', html_content, flags=re.MULTILINE)
        html_content = re.sub(r'^# (.*)$', r'<h1 style="margin: 36px 0 16px 0; font-size: 1.6em; color: #1a202c;">\1</h1>', html_content, flags=re.MULTILINE)

        # 2. Mermaid图表
        def mermaid_repl(match):
            code = match.group(1).strip()
            return (
                '<div class="mermaid-container" style="margin: 24px 0;">\n'
                '  <div class="mermaid-header" style="padding: 8px 16px; background: #f0fdf4; border-bottom: 1px solid #c6f6d5;">\n'
                '    <span class="diagram-type" style="color: #38a169; font-weight: bold;">关系图</span>\n'
                '  </div>\n'
                '  <div class="mermaid-content" style="padding: 18px; background: #f9fafb; text-align: center;">\n'
                f'    <div class="mermaid">{self._escape_html(code)}</div>\n'
                '  </div>\n'
                '</div>'
            )
        html_content = re.sub(r'```mermaid\s*\n([\s\S]*?)\n```', mermaid_repl, html_content, flags=re.DOTALL)

        # 3. 代码块
        def code_repl(match):
            lang = match.group(1) or 'text'
            code = match.group(2).strip()
            return (
                '<div class="code-block" style="margin: 24px 0;">\n'
                '  <div class="code-header" style="display: flex; align-items: center; justify-content: space-between; background: #f7fafc; padding: 8px 16px; border-bottom: 1px solid #e2e8f0;">\n'
                f'    <span class="language-tag" style="color: #3182ce; font-weight: bold;">{lang}</span>\n'
                '    <button class="copy-btn" style="background: #edf2f7; border: none; border-radius: 4px; padding: 2px 10px; cursor: pointer;">复制</button>\n'
                '  </div>\n'
                '  <pre style="margin: 0; background: #f9fafb;"><code class="language-' + lang + '">' + self._escape_html(code) + '</code></pre>\n'
                '</div>'
            )
        html_content = re.sub(r'```(\w+)?\s*\n([\s\S]*?)\n```', code_repl, html_content, flags=re.DOTALL)

        # 4. 列表美化
        html_content = re.sub(r'^\s*[-\*] (.*)$', r'<ul><li>\1</li></ul>', html_content, flags=re.MULTILINE)
        html_content = re.sub(r'^\s*\d+\. (.*)$', r'<ol><li>\1</li></ol>', html_content, flags=re.MULTILINE)
        html_content = re.sub(r'(</ul>)\s*<ul>', '', html_content)  # 合并相邻ul
        html_content = re.sub(r'(</ol>)\s*<ol>', '', html_content)

        # 5. 表格美化（简单处理）
        html_content = re.sub(r'^\|(.+?)\|$', r'<div class="table-block">|\1|</div>', html_content, flags=re.MULTILINE)

        # 6. 只对纯文本行用<p>包裹
        lines = html_content.split('\n')
        new_lines = []
        for line in lines:
            lstr = line.strip()
            if not lstr:
                continue
            # 只包裹非HTML标签开头和结尾的纯文本
            if not (lstr.startswith('<') and lstr.endswith('>')):
                new_lines.append(f'<p>{lstr}</p>')
            else:
                new_lines.append(line)
        html_content = '\n'.join(new_lines)

        return html_content
    
    def _escape_html(self, text: str) -> str:
        """转义HTML特殊字符"""
        return text.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;').replace('"', '&quot;').replace("'", '&#39;')
    
    def generate_css_styles(self) -> str:
        """生成增强样式"""
        return """
/* 多模态增强样式 */
.enhanced-table {
    width: 100%;
    border-collapse: collapse;
    margin: 15px 0;
    background: white;
    border-radius: 8px;
    overflow: hidden;
    box-shadow: 0 2px 8px rgba(0,0,0,0.1);
}

.enhanced-table th,
.enhanced-table td {
    padding: 12px 15px;
    text-align: left;
    border-bottom: 1px solid #eee;
}

.enhanced-table th {
    background: #f8f9fa;
    font-weight: 600;
    color: #2c3e50;
}

.enhanced-table tr:hover {
    background: #f8f9fa;
}

.code-block {
    margin: 15px 0;
    border-radius: 8px;
    overflow: hidden;
    box-shadow: 0 2px 8px rgba(0,0,0,0.1);
}

.code-header {
    display: flex;
    justify-content: space-between;
    align-items: center;
    padding: 8px 15px;
    background: #2d3748;
    color: white;
}

.language-tag {
    font-size: 0.9rem;
    font-weight: 500;
}

.copy-btn {
    background: #4a5568;
    color: white;
    border: none;
    padding: 4px 8px;
    border-radius: 4px;
    font-size: 0.8rem;
    cursor: pointer;
    transition: background 0.2s;
}

.copy-btn:hover {
    background: #718096;
}

.code-block pre {
    margin: 0;
    padding: 15px;
    background: #2d3748;
    color: #e2e8f0;
    overflow-x: auto;
}

.mermaid-container {
    margin: 20px 0;
    border: 1px solid #e2e8f0;
    border-radius: 8px;
    overflow: hidden;
    background: white;
}

.mermaid-header {
    padding: 10px 15px;
    background: #f7fafc;
    border-bottom: 1px solid #e2e8f0;
}

.diagram-type {
    font-size: 0.9rem;
    font-weight: 500;
    color: #4a5568;
}

.mermaid-content {
    padding: 20px;
    text-align: center;
}

/* 响应式设计 */
@media (max-width: 768px) {
    .enhanced-table {
        font-size: 0.9rem;
    }
    
    .enhanced-table th,
    .enhanced-table td {
        padding: 8px 10px;
    }
    
    .code-block pre {
        padding: 10px;
        font-size: 0.85rem;
    }
}
"""
    
    def generate_js_functions(self) -> str:
        """生成JavaScript函数"""
        return """
// 复制代码功能
function copyCode(codeId) {
    const codeBlock = document.getElementById(codeId);
    const code = codeBlock.querySelector('code').textContent;
    
    navigator.clipboard.writeText(code).then(() => {
        const btn = codeBlock.querySelector('.copy-btn');
        const originalText = btn.textContent;
        btn.textContent = '已复制!';
        btn.style.background = '#48bb78';
        
        setTimeout(() => {
            btn.textContent = originalText;
            btn.style.background = '#4a5568';
        }, 2000);
    }).catch(err => {
        console.error('复制失败:', err);
        alert('复制失败，请手动复制');
    });
}

// 初始化Mermaid图表
function initMermaid() {
    if (typeof mermaid !== 'undefined') {
        mermaid.initialize({
            startOnLoad: true,
            theme: 'default',
            flowchart: {
                useMaxWidth: true,
                htmlLabels: true
            }
        });
    }
}

// 页面加载完成后初始化
document.addEventListener('DOMContentLoaded', function() {
    initMermaid();
});
"""

# 创建全局实例
multimodal_enhancer = MultimodalEnhancer() 